#include "bzbnet/socket.h"
#include "bzbutil/platform.h"

#include "bzbio/io.h"
#include "bzbstring/stringutils.h"

#include <cstring>
#include <algorithm>
#include <memory>

#ifdef WINDOWS
#define WINDOWS_LEAN_AND_MEAN 1
#undef _WIN32_WINNT
#define _WIN32_WINNT 0x501
#include <Winsock2.h>
#include <Ws2tcpip.h>
#include <ws2tcpip.h>
#include <windows.h>

using socket_t = SOCKET;

#elif defined(LINUX)
#include <iostream>
#include <cstdio>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <unistd.h>
#include <strings.h>
#include <arpa/inet.h>

using socket_t = int;

#endif


bool sock_initialized = false;


void throw_net_error() {

#ifdef WINDOWS

	bzbutil::throw_sys_error(WSAGetLastError());

#else

	bzbutil::throw_sys_error();

#endif

}


void throw_getaddrinfo_error(int err) {

	throw std::runtime_error(
		bzbstring::concat("Socket Error: getaddrinfo(): ",
							gai_strerror(err)));
}


void deletesocket(socket_t socket_fd) {

#ifdef WINDOWS

	closesocket(socket_fd);

#elif defined(LINUX)

	::close(socket_fd);

#endif

}


void bzbnet::initialize() {

#ifdef WINDOWS

	WSADATA wsaData;
	int err = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if(err != 0) {
		bzbutil::throw_sys_error();
	}

#endif

// Linux doesn't need initialization

	sock_initialized = true;

}


void bzbnet::destroy() {

#ifdef WINDOWS

	if(sock_initialized) {
		WSACleanup();
	}

#endif

	sock_initialized = false;

}


std::vector<std::string> bzbnet::get_host_ip_list(std::string hostname) {

	if(!sock_initialized) {
		bzbnet::initialize();
	}

	addrinfo* result = nullptr;
	std::vector<std::string> address_list;

	int err = getaddrinfo(hostname.c_str(), 0, nullptr, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}

	addrinfo* curr = result;
	while(curr->ai_next) {

		if(curr->ai_family == AF_INET) {

			char buff[INET_ADDRSTRLEN];
			sockaddr_in addr_in = *((sockaddr_in*) curr->ai_addr);
			const char* res = inet_ntop(AF_INET, &addr_in.sin_addr, buff, INET_ADDRSTRLEN);
			if(!res) {
				bzbutil::throw_sys_error();
			}

			std::string IP(buff);
			if(std::count(address_list.begin(), address_list.end(), IP) == 0) {
				address_list.push_back(IP);	
			}
			
		} else if(curr->ai_family == AF_INET6) {

			char buff[INET6_ADDRSTRLEN];
			sockaddr_in6 addr_in = *((sockaddr_in6*) curr->ai_addr);
			const char* res = inet_ntop(AF_INET6, &addr_in.sin6_addr, buff, INET6_ADDRSTRLEN);
			if(!res) {
				bzbutil::throw_sys_error();
			}

			std::string IP(buff);
			if(std::count(address_list.begin(), address_list.end(), IP) == 0) {
				address_list.push_back(IP);
			}
		}

		curr = curr->ai_next;
	}

	freeaddrinfo(result);

	return address_list;
}


std::vector<std::string> bzbnet::get_host_list(std::string hostname) {

	if(!sock_initialized) {
		bzbnet::initialize();
	}

	addrinfo* result = nullptr;
	std::vector<std::string> host_list;

	int err = getaddrinfo(hostname.c_str(), 0, nullptr, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}

	if(!result) {
		bzbutil::throw_sys_error();
	}

	char host[512];
	char serv[512];

	addrinfo* curr = result;
	while(curr->ai_next) {

		int err = getnameinfo(curr->ai_addr, sizeof(sockaddr_in),
								host, 512,
								serv, 512, 0);
		// Using NI_NAMEREQD to avoid IP addresses ?
		if(err) {
			if(err == EAI_SYSTEM) {
				bzbutil::throw_sys_error();
			} else if(err != EAI_FAMILY  && err != EAI_AGAIN) {
				throw_getaddrinfo_error(err);
			}
		}
		std::string host_str(host);
		if(std::count(host_list.begin(), host_list.end(), host_str) == 0) {
			host_list.push_back(host_str);
		}

		curr = curr->ai_next;
	}

	freeaddrinfo(result);

	return host_list;
}


std::string bzbnet::get_hostname() {

	char hostname[512];
	int err = gethostname(hostname, 512);
	if(err) {
		throw_net_error();
	}

	return std::string(hostname);
}


bzbnet::TCPSocket::TCPSocket() {

	if(!sock_initialized) {
		bzbnet::initialize();
	}

	fd = nullptr;
}


bzbnet::TCPSocket::TCPSocket(void* socket_fd) {

	if(!sock_initialized) {
		throw std::runtime_error("Socket created without library initialization not permitted.");
	}

	if(*((socket_t*) socket_fd) < 0) {
		throw std::runtime_error("Socket created with negative file descriptor.");
	}

	fd = socket_fd;
}


bzbnet::TCPSocket::~TCPSocket() {

	if(fd) {
		close();
	}

}


void bzbnet::TCPSocket::connect(std::string hostname, std::string service) {

	if(!sock_initialized) {
		bzbnet::initialize();
	}

	struct addrinfo *result = nullptr, hints;

	memset(&hints, 0, sizeof(addrinfo));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;

	int err = getaddrinfo(hostname.c_str(), service.c_str(), &hints, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}

	socket_t socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
	if(!socket_fd) {
		deletesocket(socket_fd);
		throw_net_error();
	}

	err = ::connect(socket_fd, result->ai_addr, (int) result->ai_addrlen);
	freeaddrinfo(result);
	if(err < 0) {
		throw_net_error();
	}

	// Copy socket_fd to obscure variable
	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;
}


void bzbnet::TCPSocket::connect(std::string hostname, uint16_t port) {

	connect(hostname, std::to_string(port));

}


void bzbnet::TCPSocket::bind(std::string hostname, std::string service) {

	addrinfo hints, *result;
	socket_t socket_fd;

	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	getaddrinfo(hostname.c_str(), service.c_str(), &hints, &result);

	socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);

	int err = ::bind(socket_fd, result->ai_addr, result->ai_addrlen);
	if(err) {
		deletesocket(socket_fd);
		throw_net_error();
	}

	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;
}


void bzbnet::TCPSocket::bind(std::string hostname, uint16_t port) {

	bind(hostname, std::to_string(port));

}


void bzbnet::TCPSocket::bind(std::string service) {

	struct addrinfo hints, *result;
	socket_t socket_fd;

	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_flags = AI_PASSIVE;
	int err = getaddrinfo(nullptr, service.c_str(), &hints, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}


	socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);

	err = ::bind(socket_fd, result->ai_addr, result->ai_addrlen);
	if(err) {
		deletesocket(socket_fd);
		throw_net_error();
	}

	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;
}


void bzbnet::TCPSocket::bind(uint16_t port) {

	bind(std::to_string(port));

}


void bzbnet::TCPSocket::listen(std::string service, int backlog) {

	int err = ::listen(*((socket_t*) fd), backlog);
	if(err) {
		throw_net_error();
	}

}


void bzbnet::TCPSocket::listen(uint16_t port) {

	listen(std::to_string(port));

}


bzbnet::TCPSocket bzbnet::TCPSocket::accept() {

	sockaddr_storage new_addr;
	socklen_t addrsize = sizeof(sockaddr_storage);
	socket_t new_fd = ::accept(*((socket_t*) fd), (sockaddr*) &new_addr, &addrsize);
	if(addrsize < sizeof(sockaddr_storage)) {
		throw_net_error();
	}

	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &new_fd, sizeof(socket_t));
	return TCPSocket((void*) buff);
}


std::vector<char> bzbnet::TCPSocket::recv(unsigned int packetsize) {

	if(!fd || !sock_initialized) {
		throw std::runtime_error("Socket not connected. Connect and then read.");
	}

	std::vector<char> data;
	char* packet = new char[packetsize];
	unsigned int totalsize = 0;
	int size = 0;

	do {

		size = ::recv(*((socket_t*) fd), packet, packetsize, 0);

		if(size < 0) {
			throw_net_error();
		}

		if(size) {
			data.insert(data.begin() + totalsize, packet, packet + size);
			totalsize += size;
		}

	} while(size > 0);

	if(!totalsize) {

		// Connection closed
		close();
	}

	return data;
}


void bzbnet::TCPSocket::send(std::string data) {

	if(!fd || !sock_initialized || !fd) {
		throw std::runtime_error("Socket not connected. Connect and then send.");
	}

	int err = ::send(*((socket_t*) fd), &data[0], data.size(), 0);
	if(err < 0) {
		throw_net_error();
	}

}


void bzbnet::TCPSocket::send(std::vector<char> data) {

	if(!fd || !sock_initialized) {
		throw std::runtime_error("Socket not connected. Connect and then send.");
	}

	int err = ::send(*((socket_t*) fd), &data[0], data.size(), 0);
	if(err < 0) {
		throw_net_error();
	}

}


void bzbnet::TCPSocket::close() {

	if(fd) {
		deletesocket(*((socket_t*) fd));
		delete (socket_t*) fd;
		fd = nullptr;
	}
}


int bzbnet::TCPSocket::get_raw_fd() {
	return *((int*) fd);
}


bool bzbnet::TCPSocket::is_open() {
	return fd != 0;
}



bzbnet::UDPSocket::UDPSocket() {

	if(!sock_initialized) {
		bzbnet::initialize();
	}

	fd = nullptr;
}


bzbnet::UDPSocket::~UDPSocket() {

	close();
}


void bzbnet::UDPSocket::connect(std::string hostname, std::string service) {

	if(!sock_initialized) {
		bzbnet::initialize();
	}

	struct addrinfo *result = nullptr, hints;

	memset(&hints, 0, sizeof(addrinfo));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_protocol = IPPROTO_UDP;

	int err = getaddrinfo(hostname.c_str(), service.c_str(), &hints, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}

	socket_t socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
	if(!socket_fd) {
		deletesocket(socket_fd);
		throw_net_error();
	}

	err = ::connect(socket_fd, result->ai_addr, (int) result->ai_addrlen);
	freeaddrinfo(result);
	if(err < 0) {
		throw_net_error();
	}

	if(fd) {
		deletesocket(*((socket_t*) fd));
		delete (socket_t*) fd;
	}

	// Copy socket_fd to obscure variable
	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;

}


void bzbnet::UDPSocket::connect(std::string hostname, uint16_t port) {

	connect(hostname, std::to_string(port));

}


void bzbnet::UDPSocket::bind(std::string hostname, std::string service) {

	addrinfo hints, *result;
	socket_t socket_fd;

	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_protocol = IPPROTO_UDP;
	int err = getaddrinfo(hostname.c_str(), service.c_str(), &hints, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}

	socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);

	err = ::bind(socket_fd, result->ai_addr, result->ai_addrlen);
	if(err) {
		deletesocket(socket_fd);
		throw_net_error();
	}

	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;

}


void bzbnet::UDPSocket::bind(std::string hostname, uint16_t port) {

	bind(hostname, std::to_string(port));

}


void bzbnet::UDPSocket::bind(std::string service) {

	struct addrinfo hints, *result;
	socket_t socket_fd;

	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_protocol = IPPROTO_UDP;
	hints.ai_flags = AI_PASSIVE;
	int err = getaddrinfo(nullptr, service.c_str(), &hints, &result);
	if(err) {
		throw_getaddrinfo_error(err);
	}

	socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);

	err = ::bind(socket_fd, result->ai_addr, result->ai_addrlen);
	if(err) {
		deletesocket(socket_fd);
		throw_net_error();
	}

	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;
}


void bzbnet::UDPSocket::bind(uint16_t port) {

	bind(std::to_string(port));

}


void bzbnet::UDPSocket::listen(std::string service, int backlog) {

	int err = ::listen(*((socket_t*) fd), backlog);
	if(err) {
		throw_net_error();
	}

}


void bzbnet::UDPSocket::listen(uint16_t port) {

	listen(std::to_string(port));

}


std::vector<char> bzbnet::UDPSocket::recvfrom() {

	if(!sock_initialized) {
		initialize();
	}

	std::vector<char> data;
	char buff[512];

	sockaddr* addr = nullptr;
	socklen_t len = sizeof(sockaddr_storage);
	::recvfrom(*((socket_t*) fd), (void*) buff, 512, 0, addr, &len);
	data.resize(512);

	return data;
}


void bzbnet::UDPSocket::sendto(std::string data, std::string hostname, std::string service) {

	if(!sock_initialized) {
		initialize();
	}

	struct addrinfo *result = nullptr, hints;

	memset(&hints, 0, sizeof(addrinfo));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_protocol = IPPROTO_UDP;

	int err = getaddrinfo(hostname.c_str(), service.c_str(), &hints, &result);
	if(err < 0) {
		throw_getaddrinfo_error(err);
	}

	socket_t socket_fd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);

	err = ::sendto(socket_fd, (void*) &data[0], data.size(),
						0, result->ai_addr, sizeof(struct sockaddr_storage));
	if(err < 0) {
		throw_net_error();
	}

	socket_t* buff = new socket_t();
	memcpy((void*) buff, (void*) &socket_fd, sizeof(socket_t));
	fd = (void*) buff;
}


void bzbnet::UDPSocket::sendto(std::string data, std::string hostname, uint16_t port) {

	sendto(data, hostname, std::to_string(port).c_str());
}


void bzbnet::UDPSocket::close() {

	if(fd) {
		deletesocket(*((socket_t*) fd));
		delete (socket_t*) fd;
		fd = nullptr;
	}
}


int bzbnet::UDPSocket::get_raw_fd() {
	return *((int*) fd);
}

bool bzbnet::UDPSocket::is_open() {
	return fd != 0;
}
